import argparse
import multiprocessing
import os
import pickle
import sys
import warnings
import xml.etree.ElementTree as ET
from copy import deepcopy
from multiprocessing import Manager, Pipe
from threading import Thread

import cv2
import matplotlib.pyplot as plt
import numpy as np
import scipy.io as sio
import spectral.io.envi as envi
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
from numpy import flip
from spectral import imshow, view_cube
from sshkeyboard import listen_keyboard

from AutoGPU import autoGPU
from GAN_training_utils import DataResult, TrainProcess, setup_seed
# from AutoGPU import autoGPU
from models import (_1DCNN, _2DCNN, _3DCNN, _3DCNN_1DCNN, _3DCNN_AM, PURE2DCNN,
                    PURE3DCNN, PURE3DCNN_2AM, SAE, SAE_AM, DBDA_network,
                    HamidaEtAl, LeeEtAl, SSRN_network, _2dCNN, myknn, mysvm)
from myTrans2 import Generator
from NViT import ViT as NViT
from utils import DataPreProcess, myplot, plot, setpath, splitdata


class MyThread(Thread):
    def __init__(self):
        '''
        :param func: 可调用的对象
        :param args: 可调用对象的参数
        '''
        super().__init__()   # 不要忘记调用Thread的初始化方法

    def run(self):
        while True: 
            try:
                recv = parentPipe.recv()
                if recv[0] == 'plotEpoch':
                    print('='*20 + 'saving genmodel_epoch model' + '='*20)
                    torch.save(T.gen_model, resultpath + 'genmodel_epoch{}.pth'.format(recv[1]))

                if recv[0] == 'plotTest':
                    T_aux.dis_model.load_state_dict(T.bestmodel)
                    print('{}评估验证集精度为%{:.2f}对应的测试集{}'.format('*'*30, 100*recv[1], '*'*30))
                    test_acc = T_aux.evaluate(T.test_loader, T_aux.test_result)
                    print('{} set Accuracy:{:.2%}'.format('Test', test_acc))
                    print('='*20 + 'saving model' + '='*20)
                    torch.save(T_aux.dis_model, resultpath + 'bestmodel.pth')
                    myplot(processeddata, IMAGE, imagepath, T_aux.test_result,'_%{:.2f}'.format(100*test_acc))
                    print('#'*100 + 'plot successfully' + '#'*100)
            except Exception as e:
                print(e)
            
           

if __name__ == '__main__':
    result_aux = DataResult()
    print('='*30)
    setup_seed(1993)
    warnings.filterwarnings("ignore")
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str,  default='./pathology/data/032370b-20x-roi2', metavar='D',
                        help='the dataset path you load')
    parser.add_argument('--trial_number', type=int,  default=1, metavar='T',
                        help='the time you do this trial')
    parser.add_argument('--train_number', type=float,  default=0.05, metavar='NTr', 
                        help='number of training set')
    parser.add_argument('--valid_number', type=float,  default=0.01, metavar='NVa',
                        help='number of valid set')
    parser.add_argument('--test_number', type=float,  default=0.94, metavar='NTe',
                        help='number of test set')
    parser.add_argument('--patchsize', type=int,  default=9, metavar='P',
                        help='patchsize of data') 
    parser.add_argument('--modelname', type=str,  default='TransGan', metavar='P', help='which model to choose') 
    parser.add_argument('--gpu_ids', type=int,  default=-1, metavar='G',
                        help='which gpu to use')
    parser.add_argument('--gpu_num', type=int,  default=1, metavar='G',
                        help='how many gpus to use')
    parser.add_argument('--loadbestmodel', type=int,  default=0, metavar='L',
                        help='whether load model')
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 
    

    args = parser.parse_args()
    gpu_ids = args.gpu_ids
    dataset = args.dataset
    NTr = args.train_number
    trialnumber = args.trial_number
    NTe = args.test_number
    NVa = args.valid_number
    patchsize = args.patchsize
    modelname = args.modelname
    load_bestmodel = args.loadbestmodel
    gpu_num = args.gpu_num
    dark_ref = envi.open('%s.hdr'%dataset, '%s.raw'%dataset)
    dark_nparr = np.array(dark_ref.load())
    dark_nparr = flip(dark_nparr, 0)
    tree = ET.parse('%s.xml'%dataset)
    plt.figure()
    plt.imshow(dark_nparr[:,:,10])
    plt.savefig('test.jpg') 
        # get root element
    polygon = tree.getroot().findall('object/polygon')

    mask = np.zeros((dark_nparr.shape[0], dark_nparr.shape[1]), dtype="uint8")
    for p in polygon:
        x = p.findall('pt/x')
        y = p.findall('pt/y')
        x_coor = list(map(lambda x:int(x.text), x))
        y_coor = list(map(lambda y:int(y.text), y))
        c = []
        for x, y in zip(x_coor, y_coor):
            c.append([x, y])
        cor_xy = np.array(c)
        # cor_xy = np.hstack(mas(x_coor, y_coor))
        cv2.polylines(mask, np.int32([cor_xy]), 1, 1)
        cv2.fillPoly(mask, np.int32([cor_xy]), 1)


    resultpath, imagepath, datapath = setpath(dataset, trialnumber , NTr, NVa, NTe, modelname)
    IMAGE = np.flipud(dark_nparr)
    GND = mask
    spliteddata = splitdata(IMAGE, GND, datapath , trainnum=NTr, validnum=NVa, testnum=NTe)
    
    parentPipe, childPipe = Pipe(True)


    
    print('patch 的尺寸为') 
    processeddata = DataPreProcess(IMAGE, patchsize, datapath, 1).processeddata
    print(processeddata['train'].patch.shape)

    if gpu_ids != -1:
        os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_ids)
    else:
        while True:
            try:
                autoGPU(gpu_num, 4000)
                break
            except Exception as e :
                print(e)
                pass
    
    print('训练比例为 训练集{} 验证集{} 测试集{}'.format(NTr, NVa, NTe))
    print('模型为{}'.format(modelname))
    dis_model = NViT(
                depth=3,
                num_classes = 3, 
                )


    dis_model_aux = NViT(
                depth=3,
                num_classes = 3, 
                )

    
    # dis_model.load_state_dict(torch.load(resultpath + 'bestmodel.pth'))

    gen_model = Generator()
    

    if load_bestmodel:
        
        modelpara = torch.load(resultpath + 'bestmodel.pth')
        
        dis_model.load_state_dict({k.replace('module.', ''):v for k,v in modelpara.items()})
        # dis_model.load_state_dict(modelpara)
        modelpara = torch.load(resultpath + 'genmodel.pth')
        # gen_model.load_state_dict(modelpara)
        gen_model.load_state_dict({k.replace('module.', ''):v for k,v in modelpara.items()})
        print('读了最佳模型, 继续训练')


    
   
    


    if gpu_num > 1:
        dis_model = nn.DataParallel(dis_model)
        gen_model = nn.DataParallel(gen_model)

    dis_model.to('cuda')
    gen_model.to('cuda')
    dis_model_aux.to('cuda')

    T = TrainProcess({'Gen':gen_model, 'Dis':dis_model},
                    processeddata=processeddata,
                    train_config='./config.yaml',
                    pipesend=childPipe,
                     plotwrapper = [IMAGE, imagepath],
                     resultpath = resultpath)    



    dismodel_aux = deepcopy(dis_model)
    genmodel_aux = deepcopy(gen_model)
    T_aux = TrainProcess({'Gen':genmodel_aux, 'Dis':dismodel_aux})

    t = MyThread()
    t.setDaemon(True) 
    t.start()
    
    T.training_start()
   
    with open(resultpath + 'result.pkl', 'wb') as f:
        pickle.dump(T.test_result, f, pickle.HIGHEST_PROTOCOL)
    # myplot(processeddata, IMAGE, imagepath, T.test_result)
